package simple.utils.data.leveldb;

import org.fusesource.leveldbjni.JniDBFactory;
import org.iq80.leveldb.*;
import org.iq80.leveldb.impl.Iq80DBFactory;
import simple.utils.data.KeyValue;
import simple.utils.data.SeekingIterable;
import simple.utils.data.SeekingIterator;
import simple.utils.data.IterableKeyValue;
import simple.utils.processing.Progress;

import java.io.File;
import java.io.IOException;
import java.util.AbstractMap;
import java.util.Iterator;
import java.util.Map;

/**
 * Created with IntelliJ IDEA.
 * User: rehak
 * Date: 10/3/13
 * Time: 8:40 AM
 * To change this template use File | Settings | File Templates.
 */
public abstract class AbstractLevelDB<KEY, VALUE> implements AutoCloseable, IterableKeyValue<KEY, VALUE> {

    private static final int WRITE_BUFFER_SIZE = 64 * 1024;
    private static final long CACHE_SIZE = 64 * 1024;
    private static final int BLOCK_SIZE = 512 * 1024;

    public interface TypedDBIterator<KEY, VALUE> extends SeekingIterator<KEY,VALUE>, AutoCloseable {

    }

    public enum ImplType {
        JAVA, JNI
    }

    private final ImplType implType;
    private final File dir;
    private final DB db;

    public AbstractLevelDB(File dir, ImplType implType, Options options) throws IOException {
        if (options == null) {
            options = getDefaultOptions();
        }
        if (!dir.exists()) {
            dir.mkdirs();
        }
        if (!dir.exists()) {
            throw new IOException("Unable to create directory " + dir.getAbsolutePath());
        }
        this.dir = dir;
        this.implType = implType;
        this.db = open(options);
    }

    public AbstractLevelDB(File dir, ImplType implType) throws IOException {
        this(dir, implType, getDefaultOptions());
    }

    public AbstractLevelDB(File dir) throws IOException {
        this(dir, ImplType.JNI);
    }

    public ImplType getImplType() {
        return implType;
    }

    @Override
    public void delete(KEY key) throws IOException {
        db.delete(serializeKey(key));
    }

    private DB open(Options options) throws IOException {
        DB db;
        if (implType == ImplType.JAVA) {
            db = Iq80DBFactory.factory.open(dir, options);
        } else if (implType == ImplType.JNI) {
            db = JniDBFactory.factory.open(dir, options);
        } else {
            throw new IllegalArgumentException("Unknown implementation type " + implType);
        }
        return db;
    }

    @Override
    public long records() throws IOException {
        DBIterator it = db.iterator();
        long cnt = 0;
        for (it.seekToFirst(); it.hasNext(); it.next()) {
            cnt++;
        }
        it.close();
        return cnt;
    }

    protected abstract byte[] serializeKey(KEY key) throws IOException;

    protected abstract byte[] serializeValue(VALUE value) throws IOException;

    protected abstract KEY deserializeKey(byte[] key) throws IOException;

    protected abstract VALUE deserializeValue(byte[] value) throws IOException;

    @Override
    public VALUE get(KEY key) throws IOException {
        byte[] val = db.get(serializeKey(key));
        if (val != null) {
            return deserializeValue(val);
        } else {
            return null;
        }
    }

    @Override
    public void put(KEY key, VALUE value) throws IOException {
        db.put(serializeKey(key), serializeValue(value));
    }

    public void putBatch(Iterable<Map.Entry<KEY, VALUE>> data) throws IOException {
        putBatch(data, null, -1);
    }

    public void putBatch(Iterable<Map.Entry<KEY, VALUE>> data, int commitSize) throws IOException {
        putBatch(data, null, commitSize);
    }

    public void putBatch(Iterable<Map.Entry<KEY, VALUE>> data, Progress p, int commitSize) throws IOException {
        WriteOptions wro = new WriteOptions();
        wro.sync(true);
        WriteBatch batch = db.createWriteBatch();
        try {
            int written = 0;
            Iterator<Map.Entry<KEY, VALUE>> it = data.iterator();
            while (it.hasNext()) {
                Map.Entry<KEY, VALUE> e = it.next();
                byte[] key = serializeKey(e.getKey());
                byte[] value = serializeValue(e.getValue());
                batch.put(key, value);
                written++;
                if (p != null) {
                    p.addCompleted(1);
                }
                if (commitSize > 0 && written > 0 && written >= commitSize) {
                    db.write(batch, wro);
                    batch.close();
                    batch = db.createWriteBatch();
                    written = 0;
                }
            }
        } finally {
            db.write(batch, wro);
            batch.close();
        }
    }

    public File getDir() {
        return dir;
    }

    public DB getDb() {
        return db;
    }

    public static Options getDefaultOptions() {
        Options options = new Options();
        options.createIfMissing(true);
        options.writeBufferSize(WRITE_BUFFER_SIZE);
        options.cacheSize(CACHE_SIZE);
        options.blockSize(BLOCK_SIZE);
        return options;
    }

    @Override
    public void close() throws Exception {
        db.close();
    }

    public TypedDBIterator<KEY, VALUE> typedIterator() {
        final DBIterator it = dbIterator();
        it.seekToFirst();
        return new TypedDBIterator<KEY, VALUE>() {

            private volatile boolean closed = false;

            @Override
            public boolean hasNext() {
                if (closed) {
                    return false;
                }
                return it.hasNext();
            }

            @Override
            public Map.Entry<KEY, VALUE> next() {
                Map.Entry<byte[], byte[]> e = it.peekNext();
                it.next();
                try {
                    final KEY key = deserializeKey(e.getKey());
                    final VALUE val = deserializeValue(e.getValue());
                    if (!it.hasNext()) {
                        it.close();
                        closed = true;
                    }
                    return new AbstractMap.SimpleImmutableEntry<KEY, VALUE>(key, val);
                } catch (IOException e1) {
                    throw new RuntimeException(e1);
                }
            }

            @Override
            public void remove() {
                it.remove();
            }

            @Override
            public void close() throws Exception {
                if (!closed) {
                    it.close();
                }
            }

            @Override
            public void seek(KEY key) {
                try {
                    byte[] kb = serializeKey(key);
                    it.seek(kb);
                } catch (Exception e) {
                    throw new RuntimeException("Exception serializing key!", e);
                }
            }
        };
    }

    public DBIterator dbIterator() {
        return db.iterator();
    }

    public SeekingIterator<KEY,VALUE> iterator() {
        return typedIterator();
    }

}
